package cas.ihep.hep.analysis;

import cas.ihep.hep.util.EventBlock;
import cas.ihep.hep.util.MTRandom;
import cas.ihep.hep.util.MyApplication;
import cas.ihep.hep.util.NativeBlock;
import cas.ihep.spark.context.SparkContextExt;
import cas.ihep.spark.io.Packet;
import cas.ihep.spark.io.Serializer;
import cas.ihep.spark.loop.IterationServer;
import cas.ihep.spark.loop.IterativeRunnable;
import cas.ihep.spark.rdd.JIterativePairRDD;
import cas.ihep.spark.util.GlobalConfiguration;
import cas.ihep.spark.util.IterationClientKVListener;
import cas.ihep.spark.util.NamedResourceManager;
import cas.ihep.spark.util.PartitionLocator;
import cas.ihep.spark.util.function.Function2;
import cas.ihep.util.Finalizer;
import cas.ihep.util.MemoryBuffer;
import org.apache.spark.SparkConf;
import org.luaj.vm2.LuaTable;
import org.luaj.vm2.LuaValue;
import scala.Tuple2;
import scala.Tuple3;

import javax.script.*;
import java.io.*;
import java.net.URI;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Formatter;
import java.util.List;

public class LoopDemo extends MyApplication {
    private Bindings values;
    public LoopDemo(String[] args) throws IOException, ScriptException {
        super(args);
        ScriptEngineManager sem=new ScriptEngineManager();
        ScriptEngine engine=sem.getEngineByName("luaj");
        values=new SimpleBindings();
        try(Reader reader=new BufferedReader(new FileReader(args[0]))){
            engine.eval(reader,values);
        }
    }

    @Override
    public void run() throws Exception {
        srun(values);
    }

    @SuppressWarnings("Convert2Lambda")
    private static final Function2<DemoResult, DemoResult, DemoResult> aggregate = new Function2<DemoResult, DemoResult, DemoResult>() {
        @Override
        public DemoResult call(DemoResult v1, DemoResult v2) {
            long t;
            if(v1.time!=0 && v2.time!=0){
                t=Math.min(v1.time,v2.time);
            }else if(v1.time==0){
                t=v2.time;
            }else{
                t=v1.time;
            }
            return new DemoResult(v1.dt+v2.dt,v1.mc+v2.mc,
                    v1.total+v2.total,t);
        }
    };

    static class DemoResult implements Serializable {

        double dt, mc;

        long time;
        double total;

        DemoResult(double v1, double v2) {
            dt = v1;
            mc = v2;
            time=System.currentTimeMillis();
            total=0;
        }

        DemoResult(double v1,double v2,double tt){
            dt=v1;
            mc=v2;
            total=tt;
            time=System.currentTimeMillis();
        }

        DemoResult(double v1,double v2,double tt,long millis){
            dt=v1;
            mc=v2;
            total=tt;
            time=millis;
        }

        DemoResult() {
            dt = 0;
            mc = 0;
            time=0;
            total=0;
        }

        static{
            Serializer.registerHandler(DemoResult.class, new Packet.Serializer() {
                @Override
                public int length(Object data) {
                    return 16+16;
                }

                @Override
                public void serialize(Object data, DataOutput output) {
                    try{
                        DemoResult result=(DemoResult)data;
                        output.writeDouble(result.dt);
                        output.writeDouble(result.mc);
                        output.writeLong(result.time);
                        output.writeDouble(result.total);
                    }catch (IOException e){
                        e.printStackTrace();
                    }
                }

                @Override
                public Object deserialize(DataInput input) {
                    try{
                        double v1=input.readDouble();
                        double v2=input.readDouble();
                        long time=input.readLong();
                        double total=input.readDouble();
                        return new DemoResult(v1,v2,total,time);
                    }catch (IOException e){
                        e.printStackTrace();
                    }
                    return null;
                }
            });
        }
    }

    private static Object bindingsGetOrDefault(Bindings values,String key,Object defVal){
        if(values.containsKey(key)){
            return values.get(key);
        }
        return defVal;
    }

    private static class DemoParameter implements Serializable{
        double[] par;
        long time;
        DemoParameter(double[] p){
            par=p;
            time=System.currentTimeMillis();
        }
        DemoParameter(double[] p,long t){
            par=p;
            time=t;
        }
        static final Packet.Serializer serializer=new Packet.Serializer() {
            @Override
            public int length(Object o) {
                DemoParameter dp=(DemoParameter)o;
                return 4+dp.par.length*8+8;
            }

            @Override
            public void serialize(Object o, DataOutput dataOutput) {
                DemoParameter dp=(DemoParameter)o;
                try {
                    dataOutput.writeInt(dp.par.length);
                    for(double d:dp.par){
                        dataOutput.writeDouble(d);
                    }
                    dataOutput.writeLong(dp.time);
                }catch (IOException e){
                    throw new RuntimeException(e);
                }
            }

            @Override
            public Object deserialize(DataInput dataInput) {
                try{
                    int len=dataInput.readInt();
                    double[] par=new double[len];
                    for(int i=0;i<len;i++){
                        par[i]=dataInput.readDouble();
                    }
                    return new DemoParameter(par,dataInput.readLong());
                }catch (IOException e){
                    throw new RuntimeException(e);
                }
            }
        };
        static{
            Serializer.registerHandler(DemoParameter.class,serializer);
        }
    }

    private static void srun(Bindings values)throws Exception{
        SparkConf sconf=new SparkConf();
        sconf.setAppName("LoopDemo");
        LuaTable dataf=(LuaTable) values.get("data");
        LuaTable mcf=(LuaTable)values.get("mc");
        final int partitions=(Integer)values.get("slices");
        final int iterationNumber=(Integer)values.get("iterations");
        final int seed=(Integer)values.get("seed");
        MTRandom rand=new MTRandom(seed);
        GlobalConfiguration.setProperty("spark.loop.type",(String)bindingsGetOrDefault(values,"LoopType","udp"));
        GlobalConfiguration.setProperty("spark.loop.log",(String)bindingsGetOrDefault(values,"LogLevel","Info"));
        final int waitfortime=(Integer)bindingsGetOrDefault(values,"waitfor",0);
        try(Finalizer closer= Finalizer.create()){
            Tuple3<List<Tuple2<EventBlock,EventBlock>>,Long,Long> source=NewLoopPWA.parseInputData(
                    new URI(dataf.get(LuaValue.valueOf("file")).tojstring()).normalize(),
                    new URI(mcf.get(LuaValue.valueOf("file")).tojstring()).normalize(),partitions);
            SparkContextExt jsc=closer.register(new SparkContextExt(sconf));
            long start=System.nanoTime();
            JIterativePairRDD<EventBlock,EventBlock> sourcerdd= jsc.forIterative(jsc.parallelizePairs(source._1(), 50,
                    new PartitionLocator<Tuple2<EventBlock, EventBlock>>() {
                @Override
                public List<String> getPreferredLocations(List<Tuple2<EventBlock, EventBlock>> list) {
                    if(list.isEmpty()){
                        //noinspection unchecked
                        return Collections.EMPTY_LIST;
                    }
                    ArrayList<String> ret=new ArrayList<>();
                    for(Tuple2<EventBlock,EventBlock> tp:list){
                        try {
                            ret.addAll(NewLoopPWA.getUnifiedFileSystem().locations(tp._1().getPath()));
                            ret.addAll(NewLoopPWA.getUnifiedFileSystem().locations(tp._2().getPath()));
                        }catch (IOException ignored){}
                    }
                    return ret;
                }
            }));
            final long dtNumber=source._2(),mcNumber=source._3();
            long beforeTaskStart=System.nanoTime();
            //noinspection Convert2Lambda
            final IterativeRunnable<DemoParameter, DemoResult> iRun=new IterativeRunnable<DemoParameter,DemoResult>(){
                @Override
                public void run(IterationServer<DemoParameter, DemoResult> iterationServer) {
                    double[] par=new double[8];
                    double sum0=0,sum1=0;
                    double ei=0;
                    double avgSchd=(System.nanoTime()-beforeTaskStart)/1e6;
                    for(int i=0;i<iterationNumber;i++){
                        for(int j=0;j<8;j++){
                            par[j]=rand.nextDouble();
                        }
                        long start0=System.nanoTime();
                        DemoResult value=iterationServer.broadcast(new DemoParameter(par),new DemoResult(),aggregate);
                        avgSchd+=(System.nanoTime()-start0)/1e6;
                        ei+=value.total/50.0;
                        sum0+=value.dt;
                        sum1+=value.mc;
                    }
                    double cur=System.nanoTime()-start;
                    new Formatter(System.out).format("Processing %d mc events and %d data events\n" +
                            "Final result: %f %f\n"+
                            "Duration: %f\n"+
                            "Valid computing rate: %.04f\n" +
                            "Each iteration: %.02fms\n",mcNumber,dtNumber,sum0,sum1,cur/1e6,
                            ei/avgSchd,ei/iterationNumber);
                }
            };
            IterationClientKVListener<EventBlock,EventBlock,DemoParameter,DemoResult> clientListener=
                    new IterationClientKVListener<EventBlock, EventBlock, DemoParameter, DemoResult>() {
                @Override
                public void initialize(NamedResourceManager manager, EventBlock blk0, EventBlock blk1) {
                    try {
                        Serializer.registerHandler(DemoParameter.class,DemoParameter.serializer);
                        manager.put("block0",blk0.nativeData());
                        manager.put("block1",blk1.nativeData());
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                }

                @Override
                public DemoResult iterate(NamedResourceManager manager,
                                          EventBlock blk0, EventBlock blk1, DemoParameter param) {
                    long ei=System.nanoTime();
                    NativeBlock nblk0 = (NativeBlock) manager.get("block0");
                    NativeBlock nblk1 = (NativeBlock) manager.get("block1");
                    MemoryBuffer buf0=nblk0.getBuffer().mark(),buf1=nblk1.getBuffer().mark();
                    double[] ret = NewPWACalculator.compute(param.par, buf0, blk0.count(), buf1, blk1.count(),waitfortime);
                    return new DemoResult(ret[0],ret[1],(System.nanoTime()-ei)/1e6);
                }

                @Override
                public void cleanup(NamedResourceManager namedResourceManager, EventBlock eventBlock, EventBlock eventBlock2) {

                }
            };
            sourcerdd.iterate(iRun, clientListener,DemoParameter.class,new DemoResult(),aggregate);
            for(Tuple2<EventBlock,EventBlock> t:source._1()){
                t._1().close();
                t._2().close();
            }
        }
    }
}
